
library(doMC); registerDoMC()
library(random)
library(doRNG)

library(list)
library(endorse)

rm(list = ls())

load("data.RData")

isaf.data <- cbind(indivcov, data[, c(colnames(villcov), colnames(distcov), "list.y", "list.treat", "village.id")])
isaf.data <- subset(isaf.data, data$list.treat=="control" | data$list.treat == "ISAF")
isaf.data$list.treat <- ifelse(isaf.data$list.treat == "ISAF", 1, 0)

mle.estimates <- ictreg(as.formula( paste("list.y ~ ", paste(paste(colnames(indivcov), collapse = " + "), " + ", paste(colnames(villcov), collapse = " + "), " + ", paste(colnames(distcov), collapse = " + "))) ), J = 3, data = isaf.data, method = "lm", treat = "list.treat", constrained = TRUE)

set.seed(25678)
draws <- mvrnorm(n = 3, mu = coef(mle.estimates)[c(30:58, 1:29)],
                 Sigma = matrix(2, nrow = length(coef(mle.estimates)), ncol = length(coef(mle.estimates))))

n.draws <- 5000
burnin <- 500
thin <- 10

procs <- 3

set.seed(45622)

res <- foreach(i=1:procs) %dorng% {

  result <- ictregBayes(as.formula( paste("list.y ~ ", paste(paste(colnames(indivcov), collapse = " + "), " + ",
                                                             paste(colnames(villcov), collapse = " + "), " + ",
                                                             paste(colnames(distcov), collapse = " + "))) ),
                        data = isaf.data, treat = "list.treat",
                        delta.start = draws[i, 1:29], psi.start = draws[i, 30:58], burnin = burnin,
                        n.draws = n.draws, thin = thin, delta.tune = diag(.0005, 29), psi.tune = diag(.0001, 29),
                        delta.A0 = diag(c(1/9, rep(c(1/9, 1/9), 4), rep(1/9,3), 1/9, 1/9, 1/9, 1/9, 1/9,
                            rep(1/9, length(c(colnames(villcov),colnames(distcov)))))),
                        J = 3, constrained.single = "full", sensitive.model = "probit")
  return(result)
}

res.check <- list()
for(i in 1:procs)
  res.check[[i]] <- mcmc(cbind(res[[i]]$delta, res[[i]]$psi))

res.check <- mcmc.list(res.check)

R.hat <- gelman.diag(res.check)

save(res, file = paste("results-list.RData", sep = ""))

n.draws <- 30000

thin <- 100

burnin <- 0

while (max(R.hat$psrf[, 2]) > 1.1) {

  res.last <- res
  
  res <- foreach(i=1:procs) %dorng% {
    
    result <- ictregBayes(as.formula( paste("list.y ~ ", paste(paste(colnames(indivcov), collapse = " + "), " + ",
                                                             paste(colnames(villcov), collapse = " + "), " + ",
                                                             paste(colnames(distcov), collapse = " + "))) ),
                          data = isaf.data, treat = "list.treat",
                          delta.start = res.last[[i]]$delta[nrow(res.last[[i]]$delta),],
                          psi.start = res.last[[i]]$psi[nrow(res.last[[i]]$psi),],
                          burnin = 0, n.draws = n.draws, thin = thin,
                          delta.tune = diag(.0005, 29), psi.tune = diag(.0001, 29),
                          delta.A0 = diag(c(1/9, rep(c(1/9, 1/9), 4), rep(1/9,3), 1/9, 1/9, 1/9, 1/9, 1/9,
                            rep(1/9, length(c(colnames(villcov),colnames(distcov)))))),
                          J = 3, constrained.single = "full",
                          sensitive.model = "probit")
    return(result)
  }

  res.check <- list()
  for(i in 1:procs) {
    res[[i]]$delta <- rbind(res.last[[i]]$delta, res[[i]]$delta)
    res[[i]]$psi <- rbind(res.last[[i]]$psi, res[[i]]$psi)
    
    res.check[[i]] <- mcmc(cbind(res[[i]]$delta[(nrow(res[[i]]$delta)/2):nrow(res[[i]]$delta),],
                                 res[[i]]$psi[(nrow(res[[i]]$psi)/2):nrow(res[[i]]$psi),]))
  }

  res.check <- mcmc.list(res.check)
  R.hat <- gelman.diag(res.check)

  print(R.hat)

  save(res, file = paste("results-list.RData", sep = ""))

}

